/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import alloy.ast.ASTDepthFirstReplacer;
import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.DeclIter;
import alloy.ast.FormulaPara;
import alloy.ast.FormulaSeq;
import alloy.ast.Id;
import alloy.ast.Ids;
import alloy.ast.InvocablePara;
import alloy.ast.MultiplicityExpr;
import alloy.ast.Node;
import alloy.ast.Paragraph;
import alloy.ast.QualifiedName;
import alloy.ast.QuantifiedFormula;
import alloy.ast.Signature;
import alloy.ast.SumIntExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.ast.Variables;
import alloy.semantic.AlreadyMappedException;
import alloy.semantic.LocalScope;
import alloy.type.UnionType;
import alloy.util.AlloyConstants;
import alloy.util.Dbg;

/**
 * <p>Renames variables to avoid name capture during predicate/function inlining.
 * <p>The renaming proceeds as follows.  The variable
 * <code>x</code> in paragraph <code>module/para</code> is
 * renamed to <code>module/para%x</code>, if <code>x</code>
 * appears exactly once in the paragraph.  If <code>x</code> appears
 * multiple times in the paragraph, a number indicating its lexical
 * occurrence in the paragraph is appended to the name (ie., the first
 * x is x0, then x1, etc.).
 */
public final class RenameQuantifiedVarsVisitor
    extends ASTDepthFirstReplacer
    implements AlloyConstants {

    /** name of the current {@link Paragraph} */
    private QualifiedName _curParaName;

    /** a mapping from {@link String} variable names
     * which appear multiple times in the current
     * paragraph to an {@link Integer} indicating
     * how many times the variable has been observed thus
     * far.
     */
    private Map _varNameToNum;

    private void _setupPara(Paragraph para_) {
        _curParaName = para_.getName();
        // figure out which names appear multiple times
        DuplicateNameVisitor v = new DuplicateNameVisitor();
        para_.applyVisitor(v);
        _varNameToNum = v.varNameToNum;
    }

    /** @see {#_visit(Paragraph)} */
    public Object visit(Signature sig_) {
        _setupPara(sig_);
        return visit((Node) sig_);
    }

    public Object visit(FormulaPara formPara_) {
        _setupPara(formPara_);
        // handle type parameters
        Map renamingMap = new HashMap();
        String funName = null;

        if (formPara_ instanceof InvocablePara) {
            // handle args of invocable
            InvocablePara inv = (InvocablePara) formPara_;
            //Dbg.info(inv.getName().nodeString());
            Iterator declIter = inv.getArgList().getArgDecls().getDeclIter();
            while (declIter.hasNext()) {
                _handleDecl((Decl) declIter.next(), renamingMap);
            }

            // also do not rename function name if come across it
            if (formPara_ instanceof alloy.ast.Function) {
                funName = formPara_.getName().getId().nodeString();
            }
        }
        formPara_.setLocalScope(_fixLocalScope(formPara_.getLocalScope(), renamingMap, funName));
        formPara_.setFormula(
            (FormulaSeq) formPara_.getFormula().applyReturnVisitor(
                new ExprAndTypeReplaceVisitor(
                    renamingMap,
                    "original name",
                    "globally unique name")));

        // visit formula
        formPara_.setFormula(
            (FormulaSeq) formPara_.getFormula().applyReturnVisitor(this));
        return formPara_;
    }

    /**
     * fix up a local scope to accomodate some renaming
     */
    private LocalScope _fixLocalScope(LocalScope oldScope_, Map renamingMap_, String funName) {
        LocalScope newScope = new LocalScope();
        // maintain parents
        Iterator parents = oldScope_.getParents();
        while (parents.hasNext()) {
            newScope.addParent((LocalScope) parents.next());
        }
        // rebind variables using new names
        Iterator oldNameIter = oldScope_.getBoundNames();
        while (oldNameIter.hasNext()) {
            String oldName = (String) oldNameIter.next();
            String newName = null;
            if (oldName.equals("this") || oldName.equals(funName)) {
                newName = oldName;
            } else {
                Dbg.chk(renamingMap_.containsKey(oldName));
                newName = ((VariableExpr) renamingMap_.get(oldName)).nodeString();
            }
            UnionType type = oldScope_.lookupType(oldName);
            try {
                newScope.addMapping(newName, type);
            } catch (AlreadyMappedException e) {
                Dbg.fatal("shouldn't happen");
            }
        }
        return newScope;
    }

    /**
     * Create new names for all variables of a declaration
     * and add them to a renaming map
     */
    private void _handleDecl(Decl decl_, Map renamingMap_) {
        //Dbg.info(decl_.nodeString());
        // apply renaming to RHS of decl
        decl_.setExpr(
            (MultiplicityExpr) decl_.getExpr().applyReturnVisitor(
                new ExprAndTypeReplaceVisitor(
                    renamingMap_,
                    "original name",
                    "globally unique name")));
        MultiplicityExpr expr = decl_.getExpr();
        // create new names for vars of decl
        UnionType type = expr.getType();
        Variables vars = decl_.getVariables();
        for (int i = 0; i < vars.numChildren(); i++) {
            Variable var = (Variable) vars.childAt(i);
            String varName = var.nodeString();
            // don't rename "this" or "result" since they
            // can never be shadowed
            if (varName.equals("this"))
                continue;
            String newVarName = _createNewVarName(varName);
            // replace variable
            Variable newVar = new Variable(new Id(newVarName));
	    newVar.origName = varName;
            newVar.annotateTransformReplacing("renamed to be globally unique", var);
            vars.setChild(i, newVar);
	    newVar.paragraphName = (QualifiedName)_curParaName.copy();
	    
            // add to renaming map
            VariableExpr newVarExpr = new VariableExpr((Variable) newVar.copy());
            newVarExpr.setType(type);
            renamingMap_.put(varName, newVarExpr);
        }
        // do renaming for any new names created in RHS of decl
        decl_.setExpr((MultiplicityExpr) expr.applyReturnVisitor(this));
    }

    private String _createNewVarName(String varName_) {
        String newVarName = _curParaName + NAME_SEPARATOR + varName_;
        // see if var name is duplicated
        if (_varNameToNum.containsKey(varName_)) {
            Integer varInt = (Integer) _varNameToNum.get(varName_);
            newVarName += NAME_SEPARATOR + varInt.toString();
            // increment integer for variable
            _varNameToNum.put(varName_, new Integer(varInt.intValue() + 1));
        }
        return newVarName;
    }

    /**
     * Visit a var. creator node.  Rename its
     * variables, and patch up its {@link LocalScope}
     * appropriately.  Then, replace references to
     * the old variable names in its body to the new
     * names.
     */
    private Object _visit(VarCreator varCreator_) {
        Map renamingMap = new HashMap();
        Iterator declIter = varCreator_.getDecls().getDeclIter();
        while (declIter.hasNext()) {
            Decl decl = (Decl) declIter.next();
            _handleDecl(decl, renamingMap);
        }
        varCreator_.setLocalScope(
            _fixLocalScope(varCreator_.getLocalScope(), renamingMap, null));
        varCreator_.setBody(
            (Node) varCreator_.getBody().applyReturnVisitor(
                new ExprAndTypeReplaceVisitor(
                    renamingMap,
                    "original name",
                    "globally unique name")));
        varCreator_.setBody((Node) varCreator_.getBody().applyReturnVisitor(this));
        return varCreator_;
    }

    /**
     * @see #_visit(VarCreator)
     */
    public Object visit(QuantifiedFormula quantFormula_) {
        return _visit((VarCreator) quantFormula_);
    }

    /**
     * @see #_visit(VarCreator)
     */
    public Object visit(ComprehensionExpr compExpr_) {
        return _visit((VarCreator) compExpr_);
    }

    /**
     * @see #_visit(VarCreator)
     */
    public Object visit(SumIntExpr sumIntExpr_) {
        return _visit((VarCreator) sumIntExpr_);
    }

    /**
     * Detects duplicate names in a subtree of the AST.
     */
    private class DuplicateNameVisitor extends ASTDepthFirstVisitor {

        public Map varNameToNum = new HashMap();

        private Set _seenNames = new HashSet();

        private void _visit(VarCreator varCreator_) {
            DeclIter iter = new DeclIter(varCreator_);
            while (iter.hasNext()) {
                Variable var = iter.nextVar();
                String varName = var.nodeString();
                _handleVarName(varName);
            }
        }

        private void _handleVarName(String varName_) {
            if (_seenNames.contains(varName_)) { // already seen this name
                // add to map if necessary
                if (!varNameToNum.containsKey(varName_)) {
                    varNameToNum.put(varName_, new Integer(0));
                }
            } else {
                // add to seen names
                _seenNames.add(varName_);
            }
        }

        public void visit(QuantifiedFormula quantForm_) {
            this.visit((VarCreator) quantForm_);
        }

        public void visit(ComprehensionExpr compExpr_) {
            this.visit((VarCreator) compExpr_);
        }

        public void visit(SumIntExpr sumIntExpr_) {
            this.visit((VarCreator) sumIntExpr_);
        }

        /** Visit declaration.  This method should only be called
         * for visiting formals of a function.
         */
        public void visit(Decl decl_) {
            if (decl_.isEmpty())
                return;
            Iterator varIter = decl_.getVariables().getVariableIter();
            while (varIter.hasNext()) {
                String varName = ((Variable) varIter.next()).nodeString();
                _handleVarName(varName);
            }
        }

        /**
         * Visit type params of a paragraph.
         */
        public void visit(Ids typeParams_) {
            Iterator idIter = typeParams_.getIdIter();
            while (idIter.hasNext()) {
                _handleVarName(((Id) idIter.next()).nodeString());
            }
        }
    }
}
